﻿using System;
using Bettzueche.RLGlue.RLComponent;
using net = Bettzueche.RLGlue.Codec.Network;
using System.Net;
using Bettzueche.RLGlue.RLTypes;
using System.Diagnostics;

namespace Bettzueche.RLGlue.Codec {
    /// <summary>
    /// This is the network connection for an experiment program that will talk to the C/C++ rl_glue executable over sockets.
    /// </summary>
    /// <remarks>
    /// Based on Java-Glue by Brian Tanner
    /// Copyright (C) 2007, Brian Tanner
    /// http://rl-glue-ext.googlecode.com/
    /// </remarks>
    public class NetworkGlue : IExperiment {

        #region Fields

        private net.Network _network;
        private string _host = net.Network.kDefaultHost;
        private int    _port = net.Network.kDefaultPort;
        private readonly object syncLock = new object(); 

        #endregion
        
        #region Constructors

        /// <summary>
        /// Initializes a new instance of the <see cref="NetworkGlue"/> class.<para>
        /// Sets default host and port.
        /// </para>
        /// </summary>
        /// <remarks>
        /// If the user-OS-environment variables "RLGLUE_HOST" and "RLGLUE_PORT" are set, those are used.
        /// </remarks>
        public NetworkGlue() {
            String envVariableHostString = Environment.GetEnvironmentVariable("RLGLUE_HOST",EnvironmentVariableTarget.User);
            String envVariablePortString = Environment.GetEnvironmentVariable("RLGLUE_PORT",EnvironmentVariableTarget.User);
            setHostAndPorts(envVariableHostString, envVariablePortString);

        }

        /// <summary>
        /// Initializes a new instance of the <see cref="NetworkGlue"/> class.<para>
        /// Custom host, default port</para>
        /// </summary>
        /// <param name="host">The host.</param>
        public NetworkGlue(string host) {
            setHostAndPorts(host, null);
        }

        /// <summary>
        /// Initializes a new instance of the <see cref="NetworkGlue"/> class.<para>
        /// Specify custom host and port to connect to rl_glue</para>
        /// </summary>
        /// <param name="host">The host.</param>
        /// <param name="port">The port.</param>
        public NetworkGlue(string host, int port) {
            setHostAndPorts(host, "" + port);
        }

        #endregion

        #region IExperiment Members

        /// <summary>
        /// Gets number of steps in the current (or last) episode.
        /// </summary>
        public int StepCount {
            get { 
                lock(syncLock) {
                    sendEmpty(net.Protocol.kRLNumSteps, "RL_num_steps");
                    return _network.GetInt();
                }
            }
        }

        /// <summary>
        /// Gets total reward of current episode, last episode respectively.
        /// </summary>
        public double Return {
            get { 
                lock(syncLock) {
                    sendEmpty(net.Protocol.kRLReturn, "RL_return");
                    return _network.GetDouble();
                } 
            }
        }

        /// <summary>
        /// Gets the episode count.
        /// </summary>
        public int Episodes {
            get { 
                lock(syncLock) {
                    sendEmpty(net.Protocol.kRLNumEpisodes, "RL_num_episodes");
                    return _network.GetInt();
                }
            }
        }

        /// <summary>
        /// Initializes Environment and Agent.
        /// </summary>
        /// <returns>Task Spec String</returns>
        public string Init() {
            lock (syncLock) { // necessary?
                forceConnection();

                sendEmpty(net.Protocol.kRLInit, "RL_init");
                String task_spec = _network.GetString();
                return task_spec;
            }
        }

        /// <summary>
        /// Starts experiment (starts environment and agent).
        /// </summary>
        /// <param name="observation">out Parameter of initial state.</param>
        /// <param name="action">out Parameter of first action to perform.</param>
        public void Start(out RLTypes.GeneralState observation, out RLTypes.GeneralAction action) {
            lock (syncLock) { // necessary?
                sendEmpty(net.Protocol.kRLStart, "RL_start");

                observation = _network.GetObservation();
                action = _network.GetAction();
            }
        }

        /// <summary>
        /// Performs a learning step (env_step and agent_step).
        /// </summary>
        /// <param name="reward">out Parameter Reward of this step</param>
        /// <param name="nextObservation">out Parameter next State</param>
        /// <param name="nextAction">out paremeter of next action to perform.</param>
        /// <returns>
        /// true, if the agent reaches a terminal state. false otherwise
        /// </returns>
        public bool Step(out double reward, out RLTypes.GeneralState nextObservation, out RLTypes.GeneralAction nextAction) {
            lock(syncLock) {
                sendEmpty(net.Protocol.kRLStep, "RL_step");
                bool terminal = _network.GetInt() >= 1;
                reward = _network.GetDouble();
                nextObservation = _network.GetObservation();
                nextAction = _network.GetAction();
                return terminal;
            }
        }

        /// <summary>
        /// Release (dispose) resources after an experiment.
        /// </summary>
        public void CleanUp() {
            lock(syncLock) {
                sendEmpty(net.Protocol.kRLCleanup, "RL_cleanup");
            }
        }

        /// <summary>
        /// Performs the specified number of <see cref="Step">RL_steps</see>.
        /// </summary>
        /// <param name="num_steps">Number (max) of Steps. Also 0, go without step limit to a final state.</param>
        /// <returns>exit status</returns>
        public int RunEpisode(int num_steps) {
            lock(syncLock) {
                sendInt(num_steps, net.Protocol.kRLEpisode, "RL_episode");
                int exitStatus = _network.GetInt();
                return exitStatus;
            }
        }

        #endregion

        #region IEnvironment Members

        /// <summary>
        /// Not Supported.<para>
        /// Use <see cref="NetworkGlue.Init()"/> instead.</para>
        /// </summary>
        /// <returns>
        /// NotSupportedException
        /// </returns>
        public TaskSpec.ITaskSpec InitEnvironment() {
            throw new NotSupportedException("Not Supported");
        }

        /// <summary>
        /// Initializes a new episode and gets initial state.
        /// </summary>
        /// <returns>
        /// Initial state of a learning episode (corresponding observation, respectively)
        /// </returns>
        public RLTypes.GeneralState StartEnvironment() {
            lock(syncLock) {
                sendEmpty(net.Protocol.kRLEnvStart, "RL_env_start");
                return _network.GetObservation();
            }
        }

        /// <summary>
        /// Performs specified action on the environment and returns terminal flag, reward and next state.
        /// </summary>
        /// <param name="action">action to perform</param>
        /// <param name="outReward">out Parameter, reward of the performed action.</param>
        /// <param name="outObservation">out Parameter, state (observation) after the action performed.</param>
        /// <returns>
        /// true if reached a terminal state. false otherwise
        /// </returns>
        public bool StepEnvironment(RLTypes.GeneralAction action, out double outReward, out RLTypes.GeneralState outObservation) {
            lock(syncLock) {
                send_abstract_type(action, net.Protocol.kRLEnvStep, "RL_env_step");
                bool terminal = _network.GetInt() >= 1;
                outReward = _network.GetDouble();
                outObservation = _network.GetObservation();
                return terminal;
            }
        }

        /// <summary>
        /// Not Supported.<para>
        /// Use <see cref="NetworkGlue.CleanUp()"/> instead.</para>
        /// </summary>
        public void CleanupEnvironment() {
            throw new NotSupportedException();
        }

        /// <summary>
        /// Sends message to environment.
        /// </summary>
        /// <param name="message"></param>
        /// <returns>Response message, maby empty string</returns>
        public string EnvMessage(string message) {
            lock(syncLock) {
                forceConnection();
                sendString(message, net.Protocol.kRLEnvMessage, "RL_env_message");
                return _network.GetString();
            }
        }

        #endregion

        #region IAgent Members


        /// <summary>
        /// Not Supported.<para>
        /// Use <see cref="NetworkGlue.Init()"/> instead.</para>
        /// </summary>
        /// <param name="task_spec"></param>
        public void InitAgent(string task_spec) {
            throw new NotSupportedException();
        }

        /// <summary>
        /// Gets the first action to perform depending on the specified initial state.
        /// </summary>
        /// <param name="observation">Initial state</param>
        /// <returns>First action choosed by the agent.</returns>
        public RLTypes.GeneralAction StartAgent(RLTypes.GeneralState observation) {
            lock(syncLock) {
                send_abstract_type(observation, net.Protocol.kRLAgentStart, "RL_agent_start");
                return _network.GetAction();
            }
        }

        /// <summary>
        /// Learning algorithm of the agent.<para>
        /// Based on the current position and observation the agent does his learning updates and
        /// returns the following action.</para>
        /// </summary>
        /// <param name="reward">current reward (of last performed action)</param>
        /// <param name="observation">current state (as after-effect of the last action</param>
        /// <returns>Next action, the agent want to perform according to his policy.</returns>
        public RLTypes.GeneralAction StepAgent(double reward, RLTypes.GeneralState observation) {
            lock(syncLock) {
                send_reward_observation(reward, observation, net.Protocol.kRLAgentStep, "RL_agent_step");
                return _network.GetAction();
            }
        }

        /// <summary>
        /// Marks the end of an episode and gives the last reward.
        /// </summary>
        /// <param name="reward">reward of the last and final action</param>
        public void EndAgent(double reward) {
            lock(syncLock) {
                sendDouble(reward, net.Protocol.kRLAgentEnd, "RL_agent_end");
            }
        }

        /// <summary>
        /// Not Supported.<para>
        /// Use <see cref="NetworkGlue.CleanUp()"/> instead.</para>
        /// </summary>
        public void CleanupAgent() {
            throw new NotSupportedException();
        }

        /// <summary>
        /// Message thze specified message, eg. to set parameters oder ask values.
        /// </summary>
        /// <param name="message">message according to your message protocol</param>
        /// <returns>Response message, maybe empty string</returns>
        public string MessageAgent(string message) {
            lock(syncLock) {
                forceConnection();
                sendString(message, net.Protocol.kRLAgentMessage, "RL_agent_message");
                return _network.GetString();
            }
        }

        #endregion


        /// <summary>
        /// Sets the host and ports.<para>
        /// "We'll only actually set them if they seem valid.
        /// I realize it's bad that we have this copied in AgentLoader, EnvLoader, and netGlue." (BTanner)</para>
        /// </summary>
        /// <param name="hostString">The host string.</param>
        /// <param name="portString">The port string.</param>
        private void setHostAndPorts(String hostString, String portString) {
            //Now override the default or env variable port and string with these specific settings
            if (hostString != null) {
                Dns.GetHostAddresses(hostString);
                try {
                    IPAddress[] ips = Dns.GetHostAddresses(hostString);
                    if (ips.Length > 0)
                        _host = hostString;
                    else
                        System.Diagnostics.Trace.TraceWarning("Unknown host {0}. Using default localhost instead.");
                } catch (Exception e) {
                    System.Diagnostics.Trace.TraceError("Error resolving hoststring {0}: {1}",hostString,e.Message);
                }
            }

            if (portString != null) {
                try {
                    int parsedPort = Int32.Parse(portString);
                    if (parsedPort < 0 || parsedPort > 65535)
                        System.Diagnostics.Trace.TraceWarning("Could not use port you requested: {0} is not a valid port number. Using default port instead.");
                    else
                        _port = parsedPort;
                } catch (Exception e) {
                    System.Diagnostics.Trace.TraceError("Error resolving portstring {0}: {1}",portString,e.Message);
                }
            }
        }

        #region Private Helper Methods

        private void forceConnection() {
            if (_network == null) {
                string ImplementationVersion = RLGlueCore.GetImplementationVersion();
                string SpecVersion = RLGlueCore.GetSpecVersion();

                Console.WriteLine("RL-Glue .NET Experiment Codec Version: " + SpecVersion + " (" + ImplementationVersion + ")");
                Console.WriteLine("\tConnecting to " + _host + " on port " + _port + "...");

                _network = new net.Network();

                // Connect
                _network.Connect(_host, _port, net.Network.kRetryTimeout);
                Console.WriteLine("\tExperiment Codec Connected");
                _network.ClearSendBuffer();
                _network.PutInt(net.Protocol.kExperimentConnection);
                _network.PutInt(0);
                _network.FlipSendBuffer();

                try {
                    _network.Send();
                } catch (Exception ioException) {
                    Console.Error.WriteLine(ioException.StackTrace);
                    Environment.Exit(1); // Realy?!?
                }
            }
        }

        /// <summary>
        /// "This method is called by the experiment program codec after sending something.
        /// February 9 2009 :: I am making it robust to non-blocking network.recv." (BTanner)
        /// </summary>
        /// <param name="state">The state.</param>
        private void doStandardRecv(int state)/* throws IOException*/ {
            lock(syncLock) {
                _network.ClearRecvBuffer();
                try{
                int actualRecvSize = _network.Recv(8);
                int recvSizeWithoutHeader = actualRecvSize-8;

                int glueState = _network.GetInt(0);
                int dataSize = _network.GetInt(net.Protocol.kIntSize);
                int remaining = dataSize - recvSizeWithoutHeader;

                if (remaining < 0) {
                    remaining = 0;
                }
                int remainingReceived = _network.Recv(remaining);

                _network.FlipRecvBuffer();

                // Discard the header - we should have a more elegant method for doing this.
                _network.GetInt();
                _network.GetInt();

                if (glueState != state) {
                    Console.Error.WriteLine("Not synched with server. glueState = " + glueState + " but should be " + state);
                    Environment.Exit(1); // Realy?!?!
                }
                }catch(Exception e){
                    Console.Error.WriteLine(e.Message);
                    Environment.Exit(1);
                }
            }
        }

        private void doCallWithNoParams(int state)/* throws IOException*/ {
            lock(syncLock) {
                _network.ClearSendBuffer();
                _network.PutInt(state);
                _network.PutInt(0);
                _network.FlipSendBuffer();
                _network.Send();
            }
        }

        private void sendString(String theString, int theCode, String callerName) {
            lock(syncLock) {
                try {
                    _network.ClearSendBuffer();
                    _network.PutInt(theCode);
                    _network.PutInt(net.Network.SizeOf(theString));
                    _network.PutString(theString);
                    _network.FlipSendBuffer();
                    _network.Send();

                    doStandardRecv(theCode);

                } catch (ArgumentNullException nullException) {
                    Console.Error.WriteLine("You must call RL_init before calling " + callerName);
                    Console.WriteLine(nullException.StackTrace);
                    Environment.Exit(1);
                }
                catch (Exception ioException) {
                    Console.WriteLine(ioException.StackTrace);
                    Environment.Exit(1); // realy?!?!
                }
            }
        }

        private void sendEmpty(int theCode, String callerName) {
            lock(syncLock) {
                try {
                    doCallWithNoParams(theCode);
                    doStandardRecv(theCode);
                } catch (ArgumentNullException nullException) {
                    Console.Error.WriteLine("You must call RL_init before calling " + callerName);
                    Console.WriteLine(nullException.StackTrace);
                    Environment.Exit(1);
                }catch (Exception ioException) {
                    Console.WriteLine(ioException.StackTrace);
                    Environment.Exit(1);
                } 
            }
        }

        private void sendInt(int intToSend, int theCode, String callerName) {
            lock(syncLock) {
                try {
                    _network.ClearSendBuffer();
                    _network.PutInt(theCode);
                    _network.PutInt(net.Protocol.kIntSize);
                    _network.PutInt(intToSend);
                    _network.FlipSendBuffer();
                    _network.Send();

                    doStandardRecv(theCode);
                } 
                catch (ArgumentNullException nullException) {
                    Console.Error.WriteLine("You must call RL_init before calling " + callerName);
                    Console.WriteLine(nullException.StackTrace);
                    Environment.Exit(1);
                }
                catch (Exception ioException) {
                    Console.WriteLine(ioException.StackTrace);
                    Environment.Exit(1);
                } 
            }
        }

        private void sendDouble(double doubleToSend, int theCode, String callerName) {
            lock(syncLock) {
                try {
                    _network.ClearSendBuffer();
                    _network.PutInt(theCode);
                    _network.PutInt(net.Network.SizeOf(doubleToSend));
                    _network.PutDouble(doubleToSend);
                    _network.FlipSendBuffer();
                    _network.Send();

                    doStandardRecv(theCode);
                } catch (ArgumentNullException nullException) {
                    Console.Error.WriteLine("You must call RL_init before calling " + callerName);
                    Console.WriteLine(nullException.StackTrace);
                    Environment.Exit(1);
                } catch (Exception ioException) {
                    Console.WriteLine(ioException.StackTrace);
                    Environment.Exit(1);
                } 
            }
        }

        /*
         * Added by Brian Tanner to simplify the code in here.
         * @param theObject
         * @param theCode
         * @param callerName
         */
        private void send_abstract_type(IGeneralType theObject, int theCode, String callerName) {
            Debug.Assert(theObject!=null, "Someone tried to send a null object in NetGlue from: "+callerName);
            lock(syncLock) {
                try {
                    _network.ClearSendBuffer();
                    _network.PutInt(theCode);
                    _network.PutInt(net.Network.SizeOf(theObject));
                    _network.PutGeneralType(theObject);
                    _network.FlipSendBuffer();
                    _network.Send();

                    doStandardRecv(theCode);
                } catch (ArgumentNullException nullException) {
                    Console.Error.WriteLine("NULL Exception in send_abstract_type.  Assuming it's because the network is null.\n  Is it: (network==null)="+(_network==null)+"?.\nYou must call RL_init before calling " + callerName);
                    Console.WriteLine(nullException);
                    Console.WriteLine(nullException.StackTrace);
                    Environment.Exit(1);
                } catch (Exception ioException) {
                    Console.WriteLine(ioException.StackTrace);
                    Environment.Exit(1);
                } 
            }
        }

        /// <summary>
        /// Send_reward_observations the specified the reward.
        /// </summary>
        /// <param name="theReward">The reward.</param>
        /// <param name="theObservation">The observation.</param>
        /// <param name="theCode">The code.</param>
        /// <param name="callerName">Name of the caller.</param>
        private void send_reward_observation(double theReward,GeneralState theObservation, int theCode, String callerName) {
            lock(syncLock) {
                try {
                    _network.ClearSendBuffer();
                    _network.PutInt(theCode);
                    _network.PutInt(net.Network.SizeOf(theObservation) + net.Network.SizeOf(theReward));
                    _network.PutDouble(theReward);
                    _network.PutGeneralType(theObservation);
                    _network.FlipSendBuffer();
                    _network.Send();

                    doStandardRecv(theCode);
                }
                catch (ArgumentNullException nullException) {
                    Console.Error.WriteLine("You must call RL_init before calling " + callerName);
                    Console.WriteLine(nullException.StackTrace);
                    Environment.Exit(1);
                }
                catch (Exception ioException) {
                    Console.WriteLine(ioException.StackTrace);
                    Environment.Exit(1);
                } 
            }
        }

        #endregion
    }
}
